import matplotlib.pyplot as plt
from matplotlib.widgets import Button
def visualize_attention(img_batch, roi_batch, labels):
    # 获取首个样本数据
    current_sample = [0]  # 使用列表实现闭包内的可变性
    total_samples = img_batch.shape[0]

    img, roi, label = img_batch[current_sample[0]].numpy(), roi_batch[current_sample[0]].numpy(), labels[current_sample[0]].item()
    mod_names = ['Axial', 'Sagittal', 'Coronal']
    
    ax_button = plt.axes([0.8, 0.02, 0.1, 0.04])
    button = Button(ax_button, 'Next Sample')
    # 创建图形窗口
    fig, axs = plt.subplots(3, 3, figsize=(15, 15))
    plt.suptitle(f"Label: {label} | Slice: 0/{img.shape[1]-1}", fontsize=16)
    
    # 初始化切片索引
    current_slice = [0]  # 使用列表实现闭包内的可变性
    max_slice = img.shape[1] - 1
    
    # 存储图像对象
    image_plots = [[None for _ in range(3)] for _ in range(3)]
    
    # 初始化显示
    for mod in range(3):
        # 原始图像
        image_plots[0][mod] = axs[0, mod].imshow(img[mod, 0], cmap='gray')
        axs[0, mod].set_title(f"{mod_names[mod]} Image")
        
        # ROI掩膜
        image_plots[1][mod] = axs[1, mod].imshow(roi[mod, 0], cmap='Reds', alpha=0.5)
        axs[1, mod].set_title(f"{mod_names[mod]} ROI Mask")
        
        # 叠加显示
        image_plots[2][mod] = axs[2, mod].imshow(img[mod, 0], cmap='gray')
        image_plots[2][mod] = axs[2, mod].imshow(roi[mod, 0], cmap='Reds', alpha=0.3)
        axs[2, mod].set_title("Overlay")

    def update_sample(event):
        current_sample[0] = (current_sample[0] + 1) % total_samples
        new_img = img_batch[current_sample[0]].numpy()
        new_roi = roi_batch[current_sample[0]].numpy()
        new_label = labels[current_sample[0]].item()
        
        # 更新所有切片显示
        for mod in range(3):
            # 更新原始图像
            image_plots[0][mod].set_data(new_img[mod, current_slice[0]])
            # 更新ROI掩膜
            image_plots[1][mod].set_data(new_roi[mod, current_slice[0]])
            # 更新叠加显示
            axs[2, mod].images[0].set_data(new_img[mod, current_slice[0]])
            axs[2, mod].images[1].set_data(new_roi[mod, current_slice[0]])
        
        fig.suptitle(f"Sample {current_sample[0]+1}/{total_samples} | Label: {new_label}")
        fig.canvas.draw_idle()

    button.on_clicked(update_sample)
    
    def update_display():
        """更新所有模态的显示"""
        fig.suptitle(f"Label: {label} | Slice: {current_slice[0]}/{max_slice}", fontsize=16)
        for mod in range(3):
            # 更新原始图像
            image_plots[0][mod].set_data(img[mod, current_slice[0]])
            
            # 更新ROI掩膜
            image_plots[1][mod].set_data(roi[mod, current_slice[0]])
            
            # 更新叠加显示（需要重新创建以避免透明度混合问题）
            axs[2, mod].images[0].set_data(img[mod, current_slice[0]])
            axs[2, mod].images[1].set_data(roi[mod, current_slice[0]])
        
        fig.canvas.draw_idle()
    
    def on_scroll(event):
        """鼠标滚轮事件处理"""
        if event.button == 'up':
            current_slice[0] = min(current_slice[0] + 1, max_slice)
        elif event.button == 'down':
            current_slice[0] = max(current_slice[0] - 1, 0)
        update_display()
    
    # 绑定滚轮事件
    fig.canvas.mpl_connect('scroll_event', on_scroll)
    
    # 添加键盘事件支持
    def on_key(event):
        """键盘事件处理（左右方向键）"""
        if event.key == 'right':
            current_slice[0] = min(current_slice[0] + 1, max_slice)
        elif event.key == 'left':
            current_slice[0] = max(current_slice[0] - 1, 0)
        update_display()
    
    fig.canvas.mpl_connect('key_press_event', on_key)
    
    plt.tight_layout()
    plt.show()

